
import os, sys, cv2, traceback
import numpy as np
parent_path = os.path.abspath(os.path.join(__file__, *(['..'] * 3)))
sys.path.insert(0, parent_path)

import time, threading, yaml, json
from flask import Flask, request
from gevent import pywsgi, monkey, pool

from core.task.task_pipeline import TaskPipeline
from core.server import Server as BaseServer, ServerResponse
from core.utils.logger import logging

# 打补丁，使得标准库中的阻塞式I/O操作变为非阻塞式
# monkey.patch_all()

app = Flask(__name__)

def format_data_in(org_json):
    if "rtsp_origin" in org_json: # 直播流检测
        json_data = {
            "id":org_json["rtmp_push_path"],
            "tag":"task",
            "data":{
                "tag":"obj_detect",
                "data":{
                    "tag":"RTMP",
                    "data":{
                        "SteamPull_url":org_json["rtsp_origin"],
                        "SteamPush_url":org_json["rtmp_push_path"]
                    }
                }
            }
        }
    elif ("rtmp_push_path" in org_json) & ("rtsp_origin" not in org_json): # 关闭直播流
        json_data = {
            "id":org_json["rtmp_push_path"],
            "tag":"task",
            "data":{
                "tag":"obj_detect",
                "data":{
                    "tag":"pop",
                    "data":None
                }
            }
        }
    elif "images" in org_json: # FTP检测
        json_data = {
            "id":org_json["id"],
            "tag":"task",
            "data":{
                "tag":"obj_detect",
                "data":{
                    "tag":"FTP",
                    "data":{
                        "data": org_json["images"]
                        }
                }
            }
        }
    else:
        raise Exception("wrong mode in request json")
    return json_data

def powerline_format_data_out(response):
    result = {
        "id":response["id"],
        "status":"OK",
        "data":response["data"]["post_data"]
    }
    for load_path, save_path in zip(response["data"]["data"],response["data"]["ftp_upload"]):
        result["data"][load_path]["download"] = True
        result["data"][load_path]["up_load"] = "."+save_path
    return result

def objectdetect_format_data_out(response):
    data_box = response["data"]
    response_data = {}
    for img_idx, [load_path, up_path] in enumerate(zip(data_box["data"], data_box["ftp_upload"])):
        response_data[load_path] = {
            "download": True,
            "img_size": [data_box["input_data"][img_idx].shape[1], data_box["input_data"][img_idx].shape[0]],
            "object": { idx: value for idx, value in enumerate(data_box["post_data"][img_idx])},
            "up_load": up_path,
        }

    json_data = {
        "status": "OK" if response["code"] == 200 else "FAIL",
        "data": response_data,
        "id": response["id"]
    }

    
    return json_data

# DEV 旧版关闭直播回传格式
def format_data_out_rtmp_stop(data):
    json_data = {
        "status": "OK" if data["code"] == 200 else "FAIL",
        "data": {},
    }
    return json_data

def format_data_out(data):
    json_data = {
        "status": "OK" if data["code"] == 200 else "FAIL",
        "data": {
            "data": data["message"],
            "code": data["code"],
            "download": True if data["code"] == 200 else False,
            "img_size": data["img_size"],
            "object": data["data"],
            },
        "id": data["id"]
    }
    return json_data

class Server(BaseServer):
    def __init__(self, server_cfg_path: str) -> None:
        super().__init__(server_cfg_path)

# ------ 系统相关api ------

@app.route("/api/option/restart",methods=["get"])
def handle_restart():
    try:
        server.restart()
        return {"code":200}
    except Exception as e:
        logging.error(traceback.format_exc())

# ------ 功能api ------
test_flag = True

@app.route("/deepstream_rtsp_detect",methods=["post"])
def handle_rtmp_detect():
    """
        json_data = {
            "rtmp_push_path":"rtmp://218.94.155.86:1936/live/livestream_74_2",
            "rtsp_origin":"rtmp://218.94.155.86:1936/live/livestream_74_1",
            }
    """
    databox = {}
    try:
        code = 200
        message = ""
        logging.info(f"|deepstream_rtsp_detect|")
        json_data = request.json
        json_data = format_data_in(json_data)
        databox = server.main_request(data=json_data)
    except Exception as e:
        code = 500
        message = traceback.format_exc()
        logging.error(message)
    response = {"code":code,"message":message, 
                "data":databox["post_data"] if databox.get("post_data",False) else [], 
                "img_size":databox["img_size"] if databox.get("img_size",False) else [], 
                "id": databox["id"] if "id" in databox else None
               }
    response = format_data_out(response)
    logging.info("|task id {} response|".format(json_data["id"]))

    return response

@app.route("/stop_deepstream_rtsp_detect",methods=["post"])
def handle_rtmp_detect_stop():
    databox = {}
    try:
        code = 200
        message = ""
        logging.info(f"|stop_deepstream_rtsp_detect|")
        json_data = request.json
        json_data = format_data_in(json_data)
        databox = server.main_request(data=json_data)
    except Exception as e:
        code = 500
        message = traceback.format_exc()
        logging.error(message)
    response = {"code":code,"message":message, "data":[]}
    response = format_data_out_rtmp_stop(response)
    logging.info("|task id {} response|".format(json_data["id"]))
    return response

@app.route("/water_pollution_detect",methods=["post"])
def handle_waterPollution_detect():
    databox = {}
    try:
        code = 200
        message = ""
        logging.info(f"|water_pollution_detect|")
        json_data = request.json
        json_data = format_data_in(json_data)
        # 切换为水域污染任务
        json_data["data"]["data"]["tag"] = "WaterPollution_FTP"
        databox = server.main_request(data=json_data)
        json_data["data"]["data"]["tag"] = "pop"
        server.main_request(data=json_data)
    except Exception as e:
        code = 500
        message = traceback.format_exc()
        logging.error(message)
    response = {"code":code,"message":message, 
                "data":databox["post_data"] if databox.get("post_data",False) else [], 
                "img_size":databox["img_size"] if databox.get("img_size",False) else [], 
                "id": databox["id"] if "id" in databox else None
                }
    response = format_data_out(response)
    logging.info("|task id {} response|".format(json_data["id"]))
    return response

@app.route("/powerLine_detect",methods=["post"])
def handle_powerLine_detect():
    databox = {}
    try:
        code = 200
        message = ""
        logging.info(f"|powerLine_detect|")
        json_data = request.json
        json_data = format_data_in(json_data)
        # 切换为线路遮掩模型
        json_data["data"]["tag"] = "powerline_cover"
        json_data["data"]["data"]["tag"] = "FTP"
        databox = server.main_request(data=json_data)
        # json_data["data"]["data"]["tag"] = "pop"
        # server.main_request(data=json_data)
    except Exception as e:
        code = 500
        message = traceback.format_exc()
        logging.error(message)
    response = {
        "id": databox["id"],
        "code":code,
        "message":message, 
        "meta":{
            "img_size": databox["img_size"] if databox.get("img_size",False) else [], 
        },
        "data":databox
        }
    response = powerline_format_data_out(response)
    logging.info("|task id {} response|".format(json_data["id"]))
    return response

@app.route("/obj_detect",methods=["post"])
def handle_all_detect():
    """
    """
    logging.info(f"|all_detect|")
    databox = {}
    try:
        code = 200
        message = ""
        json_data = request.json
        json_data = format_data_in(json_data)
        databox = server.main_request(data=json_data)
        # json_data["data"]["data"]["tag"] = "pop"
        # server.main_request(data=json_data)
    except Exception as e:
        code = 500
        message = traceback.format_exc()
        logging.error(message)
    response = ServerResponse(
        id=databox["id"] if "id" in databox else None,
        code=code,
        message=message,
        data = databox)
    response = objectdetect_format_data_out(response)
    logging.info("task id {} response".format(json_data["id"]))
    return response

@app.route("/api",methods=["post"])
def handle_main():
    """
    data: {
        id: 1 # 流水号
        tag: task # 1级请求
        data: {
            tag: # rooftop 2级请求
            data: "b64_data..."
        }
    }
    """
    try:
        logging.info(f"|detect, json data: {request.json}|")
        json_data = request.json
        databox, message = server.sub_task(json_data)
        response = {"message":message}
        return response
    except Exception as e:
        logging.error(traceback.format_exc())


if __name__ == '__main__':

    # 日志初始化
    log_path = "core/utils/logger_samples/log_cfg.yaml"
    logging.set_config(log_path)

    # 定义后台服务类
    server_cfg_path = os.path.dirname(__file__)+"/configs/server_cfg.yml"
    server = Server(server_cfg_path)

    logging.info("|server start|")
    app.run(host="0.0.0.0",port=2682,threaded=True)

    # WSGI_server = pywsgi.WSGIServer(('0.0.0.0',2682),app,
    #                                 spawn=pool.Pool()
    #                                 )
    # WSGI_server.serve_forever()
    logging.info("|server stop|")

